import uiautomator2 as u2
import time 
import threading  
from ChatConfig import config
from BaiduChat import chat
from concurrent.futures import ThreadPoolExecutor  
from DiyAutoMessageTool import find_element,find_child_element,click_element
from AIpainting import Painter
class AutoMessage:
     # 默认联系人
     linkman=config.linkman
     # 管理员，有权前往其他联系人
     master_name_list=[config.super_master_name,]
     # 群昵称,用来识别艾特
     robot_name=None
     # 等待消息的进程限制
     wait_message_threads=30
     # 旧消息
     low_text=None
     # 是否继续等待消息
     whether_wait_new_message=True
     # 连接对象
     u2_connect=None
     # 消息字典
     message_dict={}
     # 创建一个锁 防止多个线程同时操作屏幕
     send_message_lock = threading.Lock()  
     # 创建一个锁 防止多个AI请求响应
     chat_lock = threading.Lock()  
     # 创建一个锁 防止多个文生图请求
     text_to_picture_lock=threading.Lock()
     # 消息共享的错误次数
     message_error_times=config.tolerance_error_count
     # 自定义未知异常
     class UnKnownError(Exception):
          pass
     def clear_message(self,save_message_count):
          start_message=None
          # 找出第一条消息
          for i in self.message_dict:
               if self.message_dict[i]['last_message'] is None:
                    start_message=i
                    break
          # 找不到，开摆了，直接return
          if start_message is None:
               return
          message=start_message
          next_message=self.message_dict[start_message]['next_message']
          while len(self.message_dict)>save_message_count:
               del self.message_dict[message]
               message=next_message
               next_message=self.message_dict[message]['next_message']
          self.message_dict[message]['last_message']=None
     # 发送信息
     def send_message(self,message):
          self.send_message_lock.acquire()
          try:
               # 点击输入框
               if not click_element(self.u2_connect,resourceId="com.tencent.mm:id/o4q",className="android.widget.ScrollView"):
                    print("Warning:未找到输入框")
               self.u2_connect.send_keys('【AI回复】:'+message)
               if not click_element(self.u2_connect,text="发送",resourceId="com.tencent.mm:id/bql"):
                    print("Warning:未找到发送键")
          finally:
               self.send_message_lock.release()
     # AI回复信息
     def chat_send_message(self,message,username=None):
          response=""
          if not username is None:
               response+=f"对于【{username}】的提问【{message}】:\n"
          self.chat_lock.acquire()
          try:
               response+=chat(message)
          finally:
               self.chat_lock.release()
          self.send_message(response)
     # 获取机器人在群聊的群昵称
     def get_robot_name(self):
          if click_element(self.u2_connect,resourceId="com.tencent.mm:id/fq",className="android.widget.ImageView"):
               # 滑动
               time.sleep(1)
               self.u2_connect.swipe(400, 1800, 400, 600)
               robot_name_element=find_element(self.u2_connect,resourceId="android:id/summary",className="android.widget.TextView")
               # 返回
               if not robot_name_element is None:
                    self.robot_name=robot_name_element.info['text']
               else:
                    self.robot_name="私聊中"
               click_element(self.u2_connect,resourceId="com.tencent.mm:id/actionbar_up_indicator_btn",className="android.widget.ImageView")
     # 执行命令
     def execute_command(self,username,text):
          if self.robot_name=="私聊中" or text[:len(self.robot_name)+1]=='@'+self.robot_name or text[:3]=="AI:":
          # 对话
               command=None
               if self.robot_name=="私聊中":
                    command=text
               elif text[:len(self.robot_name)+1]=='@'+self.robot_name:
                    command=text[len(self.robot_name)+2:]
               else:
                    command=text[3:]
               print(f"Chat:检测到来自【{username}】的指令:{command}")
               if "关机" == command[:2]:
                    if username == config.super_master_name:
                         self.send_message("机器人：已关闭")
                         self.whether_wait_new_message=False
                    else:
                         self.send_message("你不是管理员！！！")
               elif "好好想想你是谁" == command[:7]:
                    self.get_robot_name()
               elif "前往" == command[:2]:
                    if self.robot_name=="私聊中" or username in self.master_name_list:
                         self.enter_linkman(command[2:])
                    else:
                         self.send_message(f"{command}失败！你不是管理员！")
               elif "文生图 "==command[:4]:
                    command=command[4:]
                    if "风格 "==command[:3]:
                         command=command[3:]
                         commands=command.split()
                         if len(commands)>=2 and commands[0] in Painter.model_dict:
                              self.text_to_picture(commands[1],commands[0])
                         else:
                              self.send_message(f"文生图语法错误！请检查语法并使用【AI:风格查询】查询风格确定风格可用！")
                    else:
                         self.text_to_picture(command)
               elif "风格查询"==command[:4]:
                         self.send_message(str(list(Painter.model_dict)))

               elif "添加管理员" == command[:5]:
                    if username == config.super_master_name:
                         if command[5:] in self.master_name_list:
                              self.send_message(f"添加管理员失败！{command[5:] }本就是管理员！")
                         else:
                              self.master_name_list.append(command[5:])
                              self.send_message(f"添加管理员成功！管理员列表:{self.master_name_list}")
                    else:
                         self.send_message(f"添加管理员失败！你不是超级管理员！")
               elif "删除管理员" == command[:5]:
                    if username == config.super_master_name:
                         self.master_name_list.remove(command[5:])
                         self.send_message(f"删除管理员成功！管理员列表:{self.master_name_list}")
                    else:
                         self.send_message(f"删除管理员失败！你不是超级管理员！")
               elif "菜单" == command[:2]:
                    self.send_message(config.help_text)
               elif "字典" == command[:2]:
                    message=""
                    for i in self.message_dict:
                         message+=f"{i}\n"
                    self.send_message(f"以下信息最近获取过，会被屏蔽:\n{message}")
               else:
                    if username!="未知":
                         self.chat_send_message(command,username)
                    else:
                         self.chat_send_message(command)
     # 根据 消息布局控件 更新消息字典 并返回用户名与最新消息
     def update_message_dict(self,parent_element)->object:
               # 根据树结构判断是自己的消息则跳过
               if find_child_element(find_child_element(find_child_element(parent_element,resourceId="com.tencent.mm:id/bkj",className="android.widget.LinearLayout"),className="android.widget.LinearLayout"),className="android.widget.LinearLayout"):
                    return None
               # 找消息文本对应的控件
               child_element=find_child_element(parent_element,resourceId="com.tencent.mm:id/bkl",className="android.widget.TextView")
               if child_element is None:
                    # 说明这条信息不是文本
                    return None
               text=child_element.info['text']
               # 因为屏幕滑动错误捕获到自己的信息
               if '【AI回复】:' == text[:7]:
                    return None
               # 此时成功找到文本 重置错误次数
               self.message_error_times=config.tolerance_error_count
               # 找用户名
               username=None
               # 分为私聊和群聊两种情况
               if self.robot_name=="私聊中":
                    username=find_element(self.u2_connect,resourceId="com.tencent.mm:id/obn",className="android.widget.TextView",index=0)
               else:
                    username=find_child_element(parent_element,resourceId="com.tencent.mm:id/brc",className="android.widget.TextView")
               if not username is None:
                    username=username.info['text']
               else:
                    # username='未知'
                    #没有获取到名字,跳过
                    return None
               if (username+':'+text) not in self.message_dict:
                    if self.low_text is None:
                         self.message_dict[username+':'+text]={"last_message":None,'next_message':None}
                         # 记录这条消息
                         self.low_text=username+':'+text
                    else:
                         #上一条消息指向这条消息
                         self.message_dict[self.low_text]['next_message']=username+':'+text
                         self.message_dict[username+':'+text]={"last_message":self.low_text,'next_message':None}
                         # 记录这条消息
                         self.low_text=username+':'+text
                    return {"username":username,"text":text}
     # 寻找消息布局控件
     def find_message(self,parent_element,mode)->str:
          try:
               username=None
               text=None
               if parent_element!=None:
                    username_and_text_dict=self.update_message_dict(parent_element)
                    if username_and_text_dict is None:
                         return
                    else:
                         username=username_and_text_dict['username']
                         text=username_and_text_dict['text']
                    if mode=='WAIT':
                         print(f"Info:[{username}]:{text}")
                         self.execute_command(username,text)
                         # 进行消息清理
                         if len(self.message_dict)>=config.clear_message_count:
                              self.clear_message(config.save_message_count)
                    elif mode=='GET':
                            pass
               else:
                    self.message_error_times-=1
                    # 还能忍就再试试找该条消息
                    if self.message_error_times>1:
                         self.find_message(parent_element,mode)
                         return
                    else:
                         return
          except (u2.exceptions.UiObjectNotFoundError,KeyError) as e:
               pass
               # print(f"warning:警告+{e}")
          finally:
               pass
     # message的通用函数 默认模式等待新消息
     def general_message(self,mode='WAIT'):
          messages_element=find_element(self.u2_connect,resourceId="com.tencent.mm:id/bn1")
          if not messages_element is None:
                [self.executor.submit(self.find_message, messages_element[i],mode) for i in range(len(messages_element))]        

     # 获取消息
     def get_message(self)->dict:
          return self.general_message(mode='GET')
     # 等待新消息
     def wait_new_message(self):  
               # 如果whether_wait_new_message为True则持续运行，通过whether_wait_new_message控制程序开关
               while self.whether_wait_new_message:
                    if hasattr(self,'painter') and not getattr(self,'painter') is None:
                         # 获取操作屏幕的锁
                         self.send_message_lock.acquire()
                         try:
                              weixin_send_image=self.painter.weixin_send_image()
                         finally:
                              # 释放锁
                              self.send_message_lock.release()
                              if weixin_send_image:
                                   self.send_message(f"文生图成功！已发送{self.painter.model_style}风格的{self.painter.description}的图片!")
                                   self.painter=None
                    # 获取新消息
                    self.general_message()
               # 关闭线程池，等待所有线程结束
               self.executor.shutdown(wait=True)
     # 文生图 
     def text_to_picture(self,description,model_style='写实'):
          # 如果有其他文生图任务则取消
          self.text_to_picture_lock.acquire()
          try:
               if hasattr(self,'painter') and not getattr(self,'painter') is None:
                    print(f"Warning:文生图失败！已有{self.painter.model_style}风格的{self.painter.description}的图片正在生成!")
                    self.send_message(f"文生图失败！已有{self.painter.model_style}风格的{self.painter.description}的图片正在生成!")
                    return
               painter=Painter(self.u2_connect,description,model_style)
               if painter.painting_by_description():
                    self.painter=painter
                    print(f"Info:文生图任务提交成功！正在等待图片生成")
               else:
                    self.send_message(f"文生图失败！")
          finally:
               self.text_to_picture_lock.release()
     # 搜索联系人控件
     def search_linkman_element(self,linkman):
          # 寻找当前页面有没有测试联系人，如果没有，使用搜索
               if not click_element(self.u2_connect,text=linkman,resourceId="com.tencent.mm:id/kbq",className="android.view.View",timeout=5):
                    if not click_element(self.u2_connect,resourceId="com.tencent.mm:id/jha",className="android.widget.RelativeLayout",timeout=3):
                         raise self.UnKnownError('未找到搜索框')
                    time.sleep(1)
                    # 输入框
                    if not click_element(self.u2_connect,resourceId="com.tencent.mm:id/d98",className="android.widget.EditText",timeout=3):
                         raise self.UnKnownError('未找到输入框')
                    time.sleep(1)
                    # 输入联系人
                    self.u2_connect.send_keys(linkman)
                    time.sleep(1)
                    if not click_element(self.u2_connect,text=linkman,resourceId="com.tencent.mm:id/odf",className="android.widget.TextView",timeout=5):
                         raise self.UnKnownError('未找到联系人')
               # 卡住未进入联系人
               if not find_element(self.u2_connect,resourceId="com.tencent.mm:id/obn",className="android.widget.TextView"):
                    raise self.UnKnownError('未找到联系人')
               
     # 进入特定联系人
     def enter_linkman(self,linkman):
          # 进入联系人前先 返回 取消
          while click_element(self.u2_connect,description="返回",className="android.widget.ImageView") or click_element(self.u2_connect,text="取消",resourceId="com.tencent.mm:id/b5i",className="android.widget.TextView",timeout=5):
               time.sleep(1)
          error_times=config.tolerance_error_count
          # 清空信息
          # 防止其他进程此时添加,多删几遍
          # while self.now_threads>1:
          #      self.low_text=None
          #      self.message_dict={}
          #      print("Info:正在等待其他进程结束")
          #      time.sleep(1)
          while error_times>1:
               try:     
                    # 在界面中搜索联系人对应的控件
                    self.search_linkman_element(linkman)
                    break
               except self.UnKnownError as unknow_error:
                    if error_times>1:
                         error_times-=1
                         continue
                    elif linkman!=self.linkman:
                         self.enter_linkman(self.linkman)
                    else:
                         raise unknow_error
     # 初始化信息
     def init_message(self,linkman):
          # 获取群昵称
          self.get_robot_name()
          initial_message_times=1
          # 多次尝试获取初始消息
          while initial_message_times<=config.initial_message_count:
               self.get_message()
               print(f"Info:第{initial_message_times}次初始化旧消息:{self.message_dict}")
               initial_message_times+=1
               time.sleep(3)
          if self.robot_name is None:
               self.robot_name="私聊中"
          print("Info:初始化完成，开始等待新消息:")
          self.linkman=linkman
     def __init__(self,u2_connect):
          self.u2_connect=u2_connect
          error_times=config.tolerance_error_count
          # 创建线程池实例 限制线程的数量为wait_message_threads  
          self.executor = ThreadPoolExecutor(max_workers=self.wait_message_threads) 
          try:
               while error_times>1 and self.whether_wait_new_message:
                    # 直接重启微信
                    # u2_connect.app_start('com.tencent.mm',stop=True)
                    # 启动微信分身
                    self.u2_connect.app_stop('com.tencent.mm')
                    time.sleep(1)
                    self.u2_connect.press("home") # 模拟Home键
                    time.sleep(1)
                    self.u2_connect.press("home") # 模拟Home键
                    time.sleep(1)
                    click_element(self.u2_connect,text="微信(分身)",className="android.widget.TextView",timeout=5)
                    # 进入默认指定的联系人
                    self.enter_linkman(self.linkman)
                    # 初始化信息
                    self.init_message(self.linkman)
                    self.wait_new_message()
          finally:
               self.executor.shutdown(wait=True)

if __name__ == '__main__':
     AutoMessage(u2.connect('10.11.244.199:5555'))